import zhipuai
import json
from config import AI_CONFIG


class ZhipuBot:
    def __init__(self):
        config = AI_CONFIG['zhipu']
        self.api_key = config['api_key']
        self.top_p = config['top_p']
        self.temperature = config['temperature']
        self.default_model = config['default_model']

    def generate(self, prompt, model=None, request_id=None, top_p=None, temperature=None):
        model = model if model is not None else self.default_model
        top_p = top_p if top_p is not None else self.top_p
        temperature = temperature if temperature is not None else self.temperature
        zhipuai.api_key = self.api_key

        if request_id is not None:
            response = zhipuai.model_api.invoke(
                model=model,
                prompt=prompt,
                top_p=top_p,
                temperature=temperature
            )
        else:
            response = zhipuai.model_api.invoke(
                model=model,
                prompt=prompt,
                top_p=top_p,
                temperature=temperature,
                request_id=request_id
            )

        return response
